///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/no_transform.cpp
///
/// \brief   implementation file for natural orbital transformations
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<string>
#include<vector>

#include<formic/fqmc/no_transform.h>
#include<formic/random/random.h>
#include<formic/all_data.h>
#include<formic/archive.h>
#include<formic/exception.h>
#include<formic/numeric/numeric.h>
#include<formic/mpi/interface.h>
#include<formic/fqmc/onepdm.h>
#include<formic/fqmc/sampler.h>
#include<formic/wfn/configuration.h>
#include<formic/output/output.h>
#include<formic/numeric/matrix.h>

#include<boost/shared_ptr.hpp>

namespace formic { namespace fqmc {

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Transforms the given orbitals to be equal to those supplied and then projects them
///          out of the remaining orbitals and re-orthonormalizes the basis
///
/// \param[in]      userinp        user input options
/// \param[in,out]  ham            the hamiltonian, which holds the integrals
/// \param[in,out]  co             the coefficient object whose natural orbitals we will use
/// \param[in]      rep_orb_inds   indices of orbitals that will be replaced
/// \param[in]      rep_coeff      the new orbital coefficients for the replaced orbitals
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S>
void select_orb_transform(const formic::InputBase & userinp,
                          formic::Hamiltonian<S> & ham,
                          const std::vector<int> & rep_orb_inds,
                          const std::vector<S> & rep_coeff)
{

  // get number of atomic orbitals
  const int nao = userinp.get<int>("nao");

  // check that the orbital index array is sane
  if ( rep_orb_inds.size() > nao )
    throw formic::Exception("we have an index array with more indices than we have orbitals in formic::fqmc::select_orb_transform");

  // allocate space for the new orbitals
  std::vector<S> new_orbs(nao*nao, formic::zero(S()));

  // get number of replacements
  const int nrep = rep_orb_inds.size();

  // perform transformation on root process
  if ( formic::mpi::rank() == 0 ) {

    if ( formic::mpi::rank() == 0 ) {
      formic::of << boost::format("initial transformation matrix:") << std::endl << std::endl;
      for (int j = 0; j < nao; j++) {
        for (int i = 0; i < nrep; i++)
          formic::of << " " << formic::format_number("%14.8f", rep_coeff.at(i+j*nrep));
        formic::of << std::endl;
      }
      formic::of << std::endl;
    }

    // orthonormalize the replacement coefficients
    std::vector<S> on_rep_coeff(rep_coeff);
    {
      std::vector<S> s_neg_half(nrep*nrep, formic::zero(S()));
      for (int i = 0; i < nrep; i++)
      for (int j = 0; j < nrep; j++)
        s_neg_half.at(i+j*nrep) = formic::xdotc(nao, &rep_coeff.at(j), nrep, &rep_coeff.at(i), nrep);
      //formic::xgemm('N', 'C', nrep, nrep, nao, formic::unity(S()), &rep_coeff.at(0), nrep, &rep_coeff.at(0), nrep, formic::zero(S()), &s_neg_half.at(0), nrep);
      formic::matrix_neg_half(nrep, &s_neg_half.at(0));
      std::vector<S> temp_orbs(on_rep_coeff);
      formic::xgemm('N', 'N', nrep, nao, nrep, formic::unity(S()), &s_neg_half.at(0), nrep, &temp_orbs.at(0), nrep, formic::zero(S()), &on_rep_coeff.at(0), nrep);
    }

    // print the orbitals
    if ( formic::mpi::rank() == 0 ) {
      formic::of << boost::format("transformation matrix converted into the following unitary form:") << std::endl << std::endl;
      for (int j = 0; j < nao; j++) {
        for (int i = 0; i < nrep; i++)
          formic::of << " " << formic::format_number("%14.8f", on_rep_coeff.at(i+j*nrep));
        formic::of << std::endl;
      }
      formic::of << std::endl;
    }

    // initialize the new orbitals to be equal to the current orbitals
    for (int i = 0; i < nao; i++)
      new_orbs.at(i+i*nao) = formic::unity(S());

    // replace the selected orbitals' coefficients
    for (int i = 0; i < rep_orb_inds.size(); i++) {

      // get convenient index variables
      const int repi = rep_orb_inds.at(i);

      // check index sanity
      if ( repi < 0 || repi >= nao )
        throw formic::Exception("formic::fqmc::select_orb_transform encountered an invalid replacement orbital index of %i") % repi;

      // make the replacement
      formic::xcopy(nao, &on_rep_coeff.at(i), nrep, &new_orbs.at(repi), nao);

    }

    //if ( formic::mpi::rank() == 0 ) {
    //  formic::of << boost::format("select_orb_transform new_orbs after replacement:") << std::endl << std::endl;
    //  for (int i = 0; i < nao; i++) {
    //    for (int j = 0; j < nao; j++)
    //      formic::of << " " << formic::format_number("%14.8f", new_orbs.at(j+i*nao));
    //    formic::of << std::endl;
    //  }
    //  formic::of << std::endl;
    //}

    // get the set of unchanged orbitals
    const std::vector<int> unc_orb_inds = formic::other_indices(nao, rep_orb_inds);
    //std::set<int> rep_orb_set(rep_orb_inds.begin(), rep_orb_inds.end());
    //std::set<int> unc_orb_set;
    //for (int j = 0; j < nao; j++)
    //  if ( rep_orb_set.count(j) == 0 )
    //    unc_orb_set.insert(j);
    //std::vector<int> unc_orb_inds(unc_orb_set.begin(), unc_orb_set.end());

    // project the new orbitals out of the other orbitals
    for (int i = 0; i < rep_orb_inds.size(); i++) {
      for (int j = 0; j < unc_orb_inds.size(); j++) {
        const int uncj = unc_orb_inds.at(j);
        formic::xaxpy(nao, -on_rep_coeff.at(i+uncj*nrep), &on_rep_coeff.at(i), nrep, &new_orbs.at(uncj), nao);
      }
    }

    // compute the s_neg_half matrix
    std::vector<S> s_neg_half(nao*nao, formic::zero(S()));
    formic::xgemm('N', 'C', nao, nao, nao, formic::unity(S()), &new_orbs.at(0), nao, &new_orbs.at(0), nao, formic::zero(S()), &s_neg_half.at(0), nao);
    formic::matrix_neg_half(nao, &s_neg_half.at(0));

    // orthonormalize the orbitals
    std::vector<S> temp_orbs(new_orbs);
    formic::xgemm('N', 'N', nao, nao, nao, formic::unity(S()), &s_neg_half.at(0), nao, &temp_orbs.at(0), nao, formic::zero(S()), &new_orbs.at(0), nao);

//    // print the orbitals
//    if ( formic::mpi::rank() == 0 ) {
//      formic::of << boost::format("select_orb_transform is transforming to the following orbitals:") << std::endl << std::endl;
//      for (int i = 0; i < nao; i++) {
//        for (int j = 0; j < nao; j++)
//          formic::of << " " << formic::format_number("%14.8f", new_orbs.at(j+i*nao));
//        formic::of << std::endl;
//      }
//      formic::of << std::endl;
//    }

  }

  // perform the orbital transformation
  ham.transform_orbitals(&new_orbs.at(0));

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "Select_NO_Transform" command, which replaces a subset of the
///          orbitals with a subset of the natural orbitals and transforms the integrals
///          accordingly.
///
/// \param[in]      userinp        user input options
/// \param[in,out]  ham            the hamiltonian, which holds the integrals
/// \param[in,out]  co             the coefficient object whose natural orbitals we will use
/// \param[in]      rep_orb_inds   indices of orbitals that will be replaced
/// \param[in]      nat_orb_inds   indices of natural orbitals to replace original orbitals with
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S>
void Select_NO_Transform_command_function_detail(const formic::InputBase & userinp,
                                                 formic::Hamiltonian<S> & ham,
                                                 boost::shared_ptr<formic::CoeffObj<S> > co,
                                                 const std::vector<int> & rep_orb_inds,
                                                 const std::vector<int> & nat_orb_inds)
{

  // get number of atomic orbitals
  const int nao = userinp.get<int>("nao");

  // get the number of orbitals we are replacing
  const int nrep = rep_orb_inds.size();

  // check that the orbital index arrays are sane
  if ( rep_orb_inds.size() != nat_orb_inds.size() )
    throw formic::Exception("command \"Select_NO_Transform\" requires equal numbers of indices in its two index array arguments");
  if ( rep_orb_inds.size() > nao )
    throw formic::Exception("command \"Select_NO_Transform\" encountered an index array with more indices than we have orbitals");

  // print some info
  if ( formic::mpi::rank() == 0 ) {
    formic::of << boost::format("   ==================================================") << std::endl;
    formic::of << boost::format("   ======    Natural Orbital Transformation    ======") << std::endl;
    formic::of << boost::format("   ==================================================") << std::endl;
    formic::of << std::endl;
    formic::of << boost::format("replacing current orbitals ");
    for (int i = 0; i < nrep; i++)
      formic::of << boost::format(" %2i") % rep_orb_inds.at(i);
    formic::of << std::endl;
    formic::of << boost::format("     with natural orbitals ");
    for (int i = 0; i < nrep; i++)
      formic::of << boost::format(" %2i") % nat_orb_inds.at(i);
    formic::of << std::endl;
    formic::of << std::endl;
  }

  // create the object that will control sampling
  boost::shared_ptr<formic::fqmc::Sampler<S> > sampler( new formic::fqmc::Sampler<S>(userinp,
                                                                                     userinp.get<int>("energy_burn_len"),
                                                                                     userinp.get<int>("energy_samp_len"),
                                                                                     formic::Configuration(userinp),
                                                                                     formic::LaggedFibonacci<>(),
                                                                                     co) );

  // compute natural orbitals
  std::vector<double> anocc, bnocc;
  std::vector<S> ano, bno;
  formic::fqmc::get_nos(userinp, sampler, co, anocc, bnocc, ano, bno);

  // print the natural orbitals
  if ( formic::mpi::rank() == 0 ) {
    formic::fqmc::nat_orb_text_print(nao, "alpha", anocc, ano);
    formic::fqmc::nat_orb_text_print(nao,  "beta", bnocc, bno);
  }

  // get the natural orbitals we are using
  std::vector<S> rep_coeff(nrep * nao, formic::zero(S()));
  for (int i = 0; i < nrep; i++) {

    // get convenient index variable
    const int nati = nat_orb_inds.at(i);

    // check index sanity
    if ( nati < 0 || nati >= nao )
      throw formic::Exception("command \"Select_NO_Transform\" encountered an invalid natural orbital index of %i") % nati;

    // make the replacement
    formic::xcopy(nao, &ano.at(nati), nao, &rep_coeff.at(i), nrep);

  }

  // perform the transformation
  formic::fqmc::select_orb_transform(userinp, ham, rep_orb_inds, rep_coeff);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "Select_Orb_Transform" command, which replaces a subset of the
///          orbitals with those provided.
///
/// \param[in]      userinp        user input options
/// \param[in,out]  ham            the hamiltonian, which holds the integrals
/// \param[in,out]  co             the coefficient object whose natural orbitals we will use
/// \param[in]      rep_orb_inds   indices of orbitals that will be replaced
/// \param[in]      nat_orb_inds   indices of natural orbitals to replace original orbitals with
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S>
void Select_Orb_Transform_command_function_detail(const formic::InputBase & userinp,
                                                  formic::Hamiltonian<S> & ham,
                                                  formic::Archive & arch)
{

  // get number of atomic orbitals
  const int nao = userinp.get<int>("nao");

  // get the indices of the orbitals to replace
  std::vector<int> rep_orb_inds;
  formic::mpi::read_and_bcast(arch, rep_orb_inds, "Failed to read indices of orbitals to replace in formic::fqmc::Select_Orb_Transform_command_function_detail.");

  // read new orbital coefficients
  std::vector<S> rep_coeff;
  formic::mpi::read_and_bcast(arch, rep_coeff, "Failed to read new orbital coefficients in formic::fqmc::Select_Orb_Transform_command_function_detail.");

  // check that the orbital index array is sane
  if ( rep_orb_inds.size() > nao )
    throw formic::Exception("command \"Select_Orb_Transform\" encountered an index array with more indices than we have orbitals");

  // check the size of the replacement orbital matrix
  if ( rep_coeff.size() != rep_orb_inds.size() * nao )
    throw formic::Exception("replacement orbital coefficient array is the wrong size in formic::fqmc::Select_Orb_Transform_command_function_detail");

  // print some info
  if ( formic::mpi::rank() == 0 ) {
    formic::of << boost::format("performing an orbital transformation on orbitals");
    for (int i = 0; i < rep_orb_inds.size(); i++)
      formic::of << boost::format(" %2i") % rep_orb_inds.at(i);
    formic::of << std::endl << std::endl;
  }

  // perform the transformation
  formic::fqmc::select_orb_transform(userinp, ham, rep_orb_inds, rep_coeff);

}

}

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "Select_NO_Transform" command, which replaces a subset of the
///          orbitals with a subset of the natural orbitals and transforms the integrals
///          accordingly.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::fqmc::Select_NO_Transform_command_function(formic::Archive & arch, formic::AllData & data) {

  // read name of coefficient object whose natural orbitals will be used
  std::string co_name;
  formic::mpi::read_and_bcast(arch, co_name, "Failed to read co_name in formic::fqmc::Select_NO_Transform_command_function.");

  // read list of orbitals to replace
  std::vector<int> rep_orb_inds;
  formic::mpi::read_and_bcast(arch, rep_orb_inds, "Failed to read indices of orbitals to replace in formic::fqmc::Select_NO_Transform_command_function.");

  // read list of natural orbitals to use
  std::vector<int> nat_orb_inds;
  formic::mpi::read_and_bcast(arch, nat_orb_inds, "Failed to read indices of natural orbitals to use in formic::fqmc::Select_NO_Transform_command_function.");

  // execute the replacement and transformation
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic::fqmc::Select_NO_Transform_command_function_detail(data.userinp, data.hd, data.wd.get_co(co_name), rep_orb_inds, nat_orb_inds);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic::fqmc::Select_NO_Transform_command_function_detail(data.userinp, data.hc, data.wc.get_co(co_name), rep_orb_inds, nat_orb_inds);
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"Select_NO_Transform\"") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "Select_Orb_Transform" command, which replaces a subset of the
///          orbitals with those provided.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::fqmc::Select_Orb_Transform_command_function(formic::Archive & arch, formic::AllData & data) {

  // execute the replacement and transformation
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic::fqmc::Select_Orb_Transform_command_function_detail(data.userinp, data.hd, arch);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic::fqmc::Select_Orb_Transform_command_function_detail(data.userinp, data.hc, arch);
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"Select_Orb_Transform\"") % data.userinp.get<std::string>("scalar_type");

}
